struct Solution;

impl Solution {
    pub fn count_squares(matrix: Vec<Vec<i32>>) -> i32 {
        let row = matrix.len();
        let col = matrix[0].len();
        let mut dp = vec![vec![0; col]; row];
        let mut res = 0;

        for (i, vec) in matrix.into_iter().enumerate() {
            for (j, val) in vec.into_iter().enumerate() {
                if val > 0 && i > 0 && j > 0 {
                    dp[i][j] = i32::min(dp[i - 1][j], i32::min(dp[i][j - 1], dp[i - 1][j - 1])) + 1;
                } else {
                    dp[i][j] = val;
                }
                res += dp[i][j];
            }
        }

        res
    }
}

pub fn count_square_submatrices_with_all_ones(matrix: Vec<Vec<i32>>) -> i32 {
    Solution::count_squares(matrix)
}

fn main() {

}

#[test]
fn test() {
    let matrix =
        vec![
            vec![0, 1, 1, 1],
            vec![1, 1, 1, 1],
            vec![0, 1, 1, 1]
        ];
    let res = count_square_submatrices_with_all_ones(matrix);
    println!("{}", res);
}
